# -*- coding: utf-8 -*
import re
import subprocess

from ..common import Collector
from typing import List
import numpy as np

DISK_METRICS_NUM = 20


class DiskMetricsCollector(Collector):
    """磁盘性能指标收集器：收集 r/s, rkB/s, w/s, wkB/s r_await w_await %util """
    def __init__(self, is_test):
        super().__init__(is_test)
        self.metrics_dict = {
            "r/s": 0.00,
            "rkB/s": 0.00,
            "w/s": 0.00,
            "wkB/s": 0.00,
            "r_await": 0.00,
            "w_await": 0.00,
            "%util": 0.00,
        }

    def get_metric(self, metric: str) -> float:
        return self.metrics_dict[metric]

    def collect(self) -> List[float]:
        average: List[float] = []         # 几个磁盘的性能平均值列表
        every_line: List[List[str or float]] = []
        number_pat = re.compile(r"\d+\.\d+")
        if self.is_test:
            test_file_path = Collector.get_test_file_path("iostat.txt", __file__)
            with open(test_file_path, "r", encoding="utf-8") as f:
                lines = f.readlines()
                lines = lines[1:]

            for line in lines:
                metrics: List[str] = number_pat.findall(line)
                metrics = list(map(float, metrics))
                every_line.append(metrics)
        else:
            process_out = subprocess.run(["iostat", "-d", "-x", "1", "1", "-y"], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            out_str = str(process_out.stdout, encoding="utf-8")
            start_pos = out_str.find("Device")
            out_str = out_str[start_pos:]
            all_metrics = number_pat.findall(out_str)
            for i in range(0, len(all_metrics), DISK_METRICS_NUM):
                metrics: List[str] = all_metrics[i:i+DISK_METRICS_NUM]
                metrics = list(map(float, metrics))
                every_line.append(metrics)

        every_line_np = np.array(every_line)
        totsum = np.sum(every_line_np, axis=0)
        np.seterr(divide='ignore', invalid='ignore')
        ave_np = np.divide(totsum, len(lines)) if self.is_test else \
            np.divide(totsum, len(every_line))
        average = ave_np.tolist()
        self.stash.extend(average)
        return average
